from flask import Flask, request, redirect, url_for, render_template, jsonify
import os
import uuid
import hashlib
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 5 * 1024 * 1024 * 1024  # 5GB限制
app.config['UPLOAD_FOLDER'] = 'files/uploads'
app.config['TEMP_FOLDER'] = 'files/tmp'
app.config['SECRET_KEY'] = os.getenv('APP_SECRET', 'default-secret-key')

# 创建必要目录
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['TEMP_FOLDER'], exist_ok=True)

def validate_key(provided_key):
    """验证请求中的密钥是否匹配"""
    return provided_key == app.config['SECRET_KEY']

@app.route('/python-file-server/verify/<key>')
def verify_key(key):
    """密钥验证路由"""
    if validate_key(key):
        return redirect(url_for('upload_page',key=key))
    return "Invalid access key", 403

@app.route('/python-file-server/upload')
def upload_page():
    """文件上传页面"""
    key = request.args.get('key')
    return render_template('upload.html', key=key)

@app.route('/python-file-server/upload_chunk', methods=['POST'])
def upload_chunk():
    """处理文件分块上传"""
    # 验证Header中的密钥
    api_key = request.headers.get('X-API-Key')
    if not validate_key(api_key):
        return jsonify({"error": "Unauthorized"}), 401

    # 获取上传参数
    chunk = request.files['chunk']
    chunk_index = int(request.form['chunkIndex'])
    total_chunks = int(request.form['totalChunks'])
    file_id = request.form['fileId']
    filename = secure_filename(request.form['filename'])
    
    # 保存分块到临时目录
    chunk_dir = os.path.join(app.config['TEMP_FOLDER'], file_id)
    os.makedirs(chunk_dir, exist_ok=True)
    chunk_path = os.path.join(chunk_dir, f"{chunk_index}")
    chunk.save(chunk_path)
    
    return jsonify({"message": "Chunk uploaded", "chunk": chunk_index})

@app.route('/python-file-server/merge', methods=['POST'])
def merge_file():
    """合并文件分块"""
    # 验证Header中的密钥
    api_key = request.headers.get('X-API-Key')
    if not validate_key(api_key):
        return jsonify({"error": "Unauthorized"}), 401

    # 获取合并参数
    file_id = request.form['fileId']
    filename = secure_filename(request.form['filename'])
    
    # 合并所有分块
    chunk_dir = os.path.join(app.config['TEMP_FOLDER'], file_id)
    final_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    
    with open(final_path, 'wb') as output_file:
        for i in range(int(request.form['totalChunks'])):
            chunk_path = os.path.join(chunk_dir, str(i))
            with open(chunk_path, 'rb') as chunk_file:
                output_file.write(chunk_file.read())
            os.remove(chunk_path)
        os.rmdir(chunk_dir)
    
    return jsonify({
        "message": "File merged successfully",
        "path": final_path,
        "size": os.path.getsize(final_path)
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000)